from importlib.resources import path
import os, sys
import shutil
from tqdm import tqdm
import random
import json
import cv2
import pandas as pd
from pathlib import Path
import glob
sys.path.append('..')
from utils.json import json_to_image_one_by_one
from utils.helper import read_xml, find_nodes, change_node_text, indent, write_xml


src_path = r'/root/project/AutoRepair_T7/data/t7_training_data/t7-aa-1218-ori/all_json_recity'
# label_src_path = r'/root/project/AutoRepair_T7/data/t7_training_data/t7-aa-1218-ori/all_json_recity'

dst_path = r'/root/project/AutoRepair_T7/data/t7_training_data/t7-aa-1218-ori/all_ori_for_classification'
crop_h, crop_w = 224, 224  # crop size

# 正常code为5个字母/数字，后面可能带一个数字(第六位)
# 数字为1表示前程切割，先去掉
# 数字为2表示长线，当前方案是去掉该数字
# 数字3表示缺陷衍生残留，当前方案是去掉该数字
def label_change(label):
    if len(label) == 6 and label[-1] == '1': return 'CUT'
    if len(label) == 6 and label[-1] in ['2', '3']:
        label = label[:-1]
    return label

def json_parse(json_path):
    img_path = json_path[:-4] + 'jpg'
    img_ori = cv2.imread(img_path, 1)
    h_ori, w_ori = img_ori.shape[:2]

    for index, (label, mask) in enumerate(json_to_image_one_by_one(img_shape=img_ori.shape, 
                             json_path=json_path, toImage=False, visual=False, mask_label='box')):
        label = label_change(label)
        if label == 'CUT':
            continue
        mask[mask > 0] = 128
        mask = mask.astype('uint8')
        contours, _ = cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)

        for cnt in contours:
            # (x,y), radius = cv2.minEnclosingCircle(cnt) #外接圆
            x, y, w, h = cv2.boundingRect(cnt)  # 外接矩形 [左上角&宽高]

            x_min, x_max = x, x+w
            if w < crop_w:
                x_min = max(1, x_min - (crop_w-w)//2)
                x_max = x_min + crop_w
            
            y_min, y_max = y, y+h
            if h < crop_h:
                y_min = max(1, y-(crop_h-h)//2)
                y_max += y_min + crop_h

            img_boxed = img_ori[y_min:y_max, x_min:x_max, ...]

            boxed_out = os.path.join(dst_path, 'images', label)
            os.makedirs(boxed_out, exist_ok=True)
            # try:
            cv2.imwrite(os.path.join(boxed_out, Path(img_path).stem + '_%d.png'%(index+1)), img_boxed)
            # except:
            #     print(f'label: {label}, name: {img_path}, ',
            #         f'h: {h}, x_min:{x_min}, x_max: {x_max}, y_min: {y_min}, y_max: {y_max}')

            base_tree = read_xml("./base_example.xml")
            root = base_tree.getroot()
            anno_tree = read_xml("./anno_example.xml")
            folder_node = find_nodes(base_tree, "folder")
            filename_node = find_nodes(base_tree, "filename")
            path_node = find_nodes(base_tree, "path")
            width_node = find_nodes(base_tree, "size/width")
            height_node = find_nodes(base_tree, "size/height")
            depth_node = find_nodes(base_tree, "size/depth")
            change_node_text(folder_node, label)
            change_node_text(filename_node, str(Path(img_path).name))
            change_node_text(path_node, img_path)
            change_node_text(width_node, str(img_ori.shape[0]))
            change_node_text(height_node, str(img_ori.shape[1]))
            change_node_text(depth_node, str(img_ori.shape[2]))

            xmin_node = find_nodes(anno_tree, "bndbox/xmin")
            ymin_node = find_nodes(anno_tree, "bndbox/ymin")
            xmax_node = find_nodes(anno_tree, "bndbox/xmax")
            ymax_node = find_nodes(anno_tree, "bndbox/ymax")
            change_node_text(xmin_node, str(x_min))
            change_node_text(ymin_node, str(y_min))
            change_node_text(xmax_node, str(x_max))
            change_node_text(ymax_node, str(y_max))
            root.append(anno_tree.getroot())
            indent(root)

            label_out = os.path.join(dst_path, 'annotations', label)
            os.makedirs(label_out, exist_ok=True)
            write_xml(base_tree, os.path.join(label_out, Path(img_path).stem + '_%d.xml'%(index+1)))


def main():
    

    df = pd.DataFrame(data={'json':glob.glob(os.path.join(src_path, '*/*.json'))})  # 获取路径下所有json文件
    df = df.sample(frac=1, random_state=10).reset_index(drop=True)  # 打乱DataFrame

    try:
        from pandarallel import pandarallel
        pandarallel.initialize(progress_bar=True) 
        print('Use multi threading !')
        is_pandarallel = True
    except:
        print('Use single threading !')
        is_pandarallel = False
    
    if is_pandarallel:
        df['json'].parallel_apply(lambda x: json_parse(x))
    else:
        df['json'].apply(lambda x: json_parse(x))
    print()



if __name__=='__main__':
    main()